#include <chrono>
#include <functional>
#include <memory>
#include <string>

#include "rclcpp/rclcpp.hpp"
#include "common_lib/glog/glog.h"
#include <opencv2/opencv.hpp>
#include <Eigen/Eigen>

// g2o 
#include "g2o/core/base_vertex.h"
#include "g2o/core/base_unary_edge.h"
#include "g2o/core/sparse_optimizer.h"
#include "g2o/core/block_solver.h"
#include "g2o/core/optimization_algorithm_levenberg.h"
#include "g2o/core/optimization_algorithm_gauss_newton.h"
#include "g2o/core/optimization_algorithm_dogleg.h"
#include "g2o/solvers/dense/linear_solver_dense.h"

using namespace std;

// 参考 《视觉SLAM十四讲》第二版 6.3.3 章节
// ros2 run common_test test_fitting_curve_g2o

// xxty@xxty-virtual-machine:~/trobot/ros2_ws$ ros2 run common_test test_fitting_curve_g2o
// start optimization
// iteration= 0     chi2= 376785.128234     time= 5.8216e-05        cumTime= 5.8216e-05     edges= 100      schur= 0
// iteration= 1     chi2= 35673.566018      time= 3.6321e-05        cumTime= 9.4537e-05     edges= 100      schur= 0
// iteration= 2     chi2= 2195.012304       time= 2.0495e-05        cumTime= 0.000115032    edges= 100      schur= 0
// iteration= 3     chi2= 174.853126        time= 1.5336e-05        cumTime= 0.000130368    edges= 100      schur= 0
// iteration= 4     chi2= 102.779695        time= 1.4514e-05        cumTime= 0.000144882    edges= 100      schur= 0
// iteration= 5     chi2= 101.937194        time= 1.6171e-05        cumTime= 0.000161053    edges= 100      schur= 0
// iteration= 6     chi2= 101.937020        time= 1.614e-05         cumTime= 0.000177193    edges= 100      schur= 0
// iteration= 7     chi2= 101.937020        time= 2.03e-05  cumTime= 0.000197493    edges= 100      schur= 0
// iteration= 8     chi2= 101.937020        time= 2.0073e-05        cumTime= 0.000217566    edges= 100      schur= 0
// iteration= 9     chi2= 101.937020        time= 2.0891e-05        cumTime= 0.000238457    edges= 100      schur= 0
// solve time cost = 0.00313573 seconds. 
// estimated a=0.890912 b=2.1719 c=0.943629

// 曲线模型顶点  模板参数：优化变量和数据类型
class CurveFittingVertex : public g2o::BaseVertex<3, Eigen::Vector3d>{
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

        // 重置
        virtual void setToOriginImpl() override {
            _estimate << 0, 0, 0;
        }

        // 更新顶点
        virtual void oplusImpl(const double* update) override {
            _estimate += Eigen::Vector3d(update);
        }

        // 存盘和读盘 留空
        virtual bool read(istream& in) override {}
        virtual bool write(ostream& out) const override {}
};

// 误差模型 模板参数：观测值维度、类型、连接顶点类型
class CurveFittingEdge : public g2o::BaseUnaryEdge<1, double, CurveFittingVertex>{
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

        CurveFittingEdge(double x): BaseUnaryEdge(), _x(x){}

        // 计算曲线误差模型
        virtual void computeError() override {
            const CurveFittingVertex* v = static_cast<const CurveFittingVertex*>(_vertices[0]);
            const Eigen::Vector3d abc = v->estimate();
            _error(0, 0) = _measurement - exp(abc(0,0)*_x*_x + abc(1,0)*_x + abc(2,0));
        };

        // 计算雅可比矩阵
        virtual void linearizeOplus() override {
            const CurveFittingVertex* v = static_cast<const CurveFittingVertex*>(_vertices[0]);
            const Eigen::Vector3d abc = v->estimate();
            double y = exp(abc[0]*_x*_x + abc[1]*_x + abc[2]);
            _jacobianOplusXi[0] = -_x*_x*y;
            _jacobianOplusXi[1] = -_x*y;
            _jacobianOplusXi[2] = -y;
        }

        virtual bool read(istream& in) override {}
        virtual bool write(ostream& out) const override {}
    
    public:
        double _x; //x值 x值为_meaurement
};

class FittingCurveG2O : public rclcpp::Node{
    public:
        FittingCurveG2O();
        ~FittingCurveG2O() = default;

        // 拟合曲线
        void FittingCurve();
};


FittingCurveG2O::FittingCurveG2O(): Node("test_fitting_curve_g2o"){
    FittingCurve();
}

void FittingCurveG2O::FittingCurve(){
    double ar=1.0, br=2.0, cr=1.0;      //真实参数值
    double ae=2.0, be=-1.0, ce=5.0;     //估计参数值
    int N=100;  //数据点个数
    double w_sigma=1.0;  //噪声Sigma值
    // double inv_sigma=1.0/w_sigma;
    cv::RNG rng;        //opencv 随机数产生器

    // 构造数据集
    vector<double> x_data, y_data;
    for(int i=0; i<N; i++){
        double x = i/100.0;
        x_data.push_back(x);
        y_data.push_back(exp(ar*x*x + br*x + cr) + rng.gaussian(w_sigma*w_sigma));
    }

    // 构建图优化
    // 每个误差项 优化变量维度为3 误差值维度为1
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<3, 1>> BlockSolverType; 
    // 线性求解器类型
    typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;

    // 梯度下降法，可以选取GN、LM、DogLeg
    auto solver = new g2o::OptimizationAlgorithmGaussNewton(
        g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>())
    );
    g2o::SparseOptimizer optimizer; //图模型
    optimizer.setAlgorithm(solver); //设置求解器
    optimizer.setVerbose(true);     // 打开调试输出

    // 往图中添加顶点
    CurveFittingVertex* v = new CurveFittingVertex();   //顶点
    v->setEstimate(Eigen::Vector3d(ae, be, ce));        //设置初始值
    v->setId(0);
    optimizer.addVertex(v);                             //添加顶点

    // 往图中添加边
    for(int i=0; i<N; i++){
        CurveFittingEdge* edge = new CurveFittingEdge(x_data[i]);
        edge->setId(i);
        edge->setVertex(0, v);              //设置连接的顶点
        edge->setMeasurement(y_data[i]);    //设置观测值
        //设置信息矩阵: 协方差矩阵之逆
        edge->setInformation(Eigen::Matrix<double, 1, 1>::Identity()*1/(w_sigma*w_sigma)); 
        optimizer.addEdge(edge); //添加边
    }

    // 执行优化
    cout << "start optimization" << endl;
    chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
    optimizer.initializeOptimization();
    optimizer.optimize(10);
    chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
    chrono::duration<double> time_used = chrono::duration_cast<chrono::duration<double>>(t2-t1);
    cout << "solve time cost = " << time_used.count() << " seconds. " << endl;

    // 输出优化值
    Eigen::Vector3d abc_estimate = v->estimate();
    cout << "estimated a=" << abc_estimate[0] << " b=" << abc_estimate[1] << " c=" << abc_estimate[2] << endl;
}


int main(int argc, char * argv[]){
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<FittingCurveG2O>());
    rclcpp::shutdown();
    return 0;
}